#Todo: 实现对图数据的mask和补全操作

import torch

import main_data_generate
import main_data_train
import warnings
warnings.filterwarnings("ignore")

if __name__ == '__main__':
    # 设置批量参数
    num_batches = 30  # 生成100个数据集
    num_nodes = 64  # 每个图64个节点
    zoom_levels = 0 # 变焦级别，值越大 聚合程度越大
    edges_per_new_node = 1
    train_flag = True  # 进行预测
    completion_flag = False #进行补全


    if completion_flag:
        config = {
            'batch_size': 32,
            'epochs': 512,
            'lr': 0.003,
            'hidden_dim': 128,
            'num_layers': 10,
            'dropout': 0.3,
            'task_type': 'completion',
            'device': 'cuda' if torch.cuda.is_available() else 'cpu'
        }
        main_data_generate.mask_generator(num_batches, num_nodes, edges_per_new_node, zoom_levels) # 实现方法 生成mask后的图数据
        main_data_train.train_completion(config,visual_flag=False) 


    if train_flag:

        config = {
            'batch_size': 32,
            'epochs': 512,
            'lr': 0.003,
            'hidden_dim': 128,
            'num_layers': 10,
            'dropout': 0.3,
            'task_type': 'classification',
            'device': 'cuda' if torch.cuda.is_available() else 'cpu',

            # --- DGI超参数 ---
            'dgi_hidden_dim': 128,  # DGI编码器输出的嵌入维度
            'dgi_epochs': 200,      # DGI无监督预训练的轮数
            'dgi_lr': 0.001,        # DGI预训练的学习率
        }

        # 定义一个切换开关，用于选择图生成模型
        selected_model = 'BA'  # 可选值：'BA', 'ER', 'WS', 'REGULAR', 'RGEOM'

        # 定义每种模型的参数
        model_params = {
            'BA': {
                'model': 'BA'
            },
            'ER': {
                'model': 'ER',
                'edge_prob': 0.05  # ER 模型特有参数
            },
            'WS': {
                'model': 'WS',
                'neighbors': 4,    # WS 模型特有参数
                'rewire_prob': 0.1
            },
            'REGULAR': {
                'model': 'REGULAR',
                'degree': 3        # REGULAR 模型特有参数
            },
            'RGEOM': {
                'model': 'RGEOM',
                'radius': 0.2      # RGEOM 模型特有参数
            }
        }

        # 根据切换开关选择模型参数
        if selected_model in model_params:
            print(f"Using {selected_model} model for graph generation...")
            main_data_generate.generator(
                num_batches,
                num_nodes,
                edges_per_new_node,
                zoom_levels,
                **model_params[selected_model]  # 解包对应模型的参数
            )
        else:
            print(f"Error: Unsupported model '{selected_model}'. Please choose from {list(model_params.keys())}.")
        
        # 定义一个开关，用于选择训练模型
        selected_training = 'train_graphmae2'  # 可选值：'train_gnn', 'train_graphsage', 'train_dgi', 'train_graphcl', 'train_graphmae2'
        
        # 根据开关调用对应的训练函数
        if selected_training == 'train_gnn':
            print("Running GNNModel training...")
            main_data_train.train_gnn(config, visual_flag=True)
        
        elif selected_training == 'train_graphsage':
            print("Running GraphSAGE training...")
            main_data_train.train_graphsage(config, visual_flag=True)
        
        elif selected_training == 'train_dgi':
            print("Running DGI training and evaluation...")
            main_data_train.train_dgi_and_evaluate(config, visual_flag=True)
        
        elif selected_training == 'train_graphcl':
            print("Running GraphCL training...")
            main_data_train.train_graphcl(config, visual_flag=True)
        
        elif selected_training == 'train_graphmae2':
            print("Running GraphMAE2 training...")
            main_data_train.train_graphmae2(config, visual_flag=True)
        
        else:
            print(f"Error: Unsupported training type '{selected_training}'. Please choose from "
                  f"'train_gnn', 'train_graphsage', 'train_dgi', 'train_graphcl', 'train_graphmae2'.")


